package com.bw.gmall.realtime.utils;

import com.alibaba.fastjson.JSONObject;
import com.bw.gmall.realtime.common.GmallConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import org.apache.hadoop.hbase.client.AsyncConnection;

import java.util.Collections;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


@Slf4j
public abstract class AsyncDimFunction<T> extends RichAsyncFunction<T, T> implements DimFunction<T> {

    private AsyncConnection hBaseAsyncConn;
    ThreadPoolExecutor threadPool = null;

    @Override
    public void open(Configuration parameters) throws Exception {
        hBaseAsyncConn = HBaseApiUtil.getHBaseAsyncConnection();
        threadPool = ThreadPoolUtil.getThreadPoolExecutor();
    }

    @Override
    public void close() throws Exception {
        HBaseApiUtil.closeAsyncHbaseConnection(hBaseAsyncConn);
    }

    @Override
    public void asyncInvoke(T input, ResultFuture<T> resultFuture) throws Exception {
        // Check if thread pool is available
        if (threadPool == null || threadPool.isShutdown()) {
            // Handle error case
            resultFuture.complete(Collections.singleton(input));
            return;
        }

        threadPool.submit(new Runnable() {
            @Override
            public void run() {
                try {
                    // 添加超时控制，避免HBase操作长时间阻塞
                    JSONObject jsonObject = HBaseApiUtil.readDimAsync(hBaseAsyncConn,
                            GmallConfig.HBASE_SCHEMA, //namespace
                            getTableName(),  // 表名
                            getRowKey(input));// rowkey 主键
                    
                    // 只有在获取到数据时才添加维度信息
                    if (jsonObject != null && !jsonObject.isEmpty()) {
                        addDims(input, jsonObject);
                    }
                    
                    resultFuture.complete(Collections.singleton(input));
                } catch (Exception e) {
                    log.error("Error while fetching dimension data for input: " + input, e);
                    // Complete the future even in case of error to prevent timeouts
                    resultFuture.complete(Collections.singleton(input));
                }
            }
        });
    }

    @Override
    public void timeout(T input, ResultFuture<T> resultFuture) throws Exception {
        log.warn("Async function timed out for input: " + input);
        // Even on timeout, complete the future to prevent the job from failing
        resultFuture.complete(Collections.singleton(input));
    }

}